function [P,Q,W,X,Y,Z] = PTF2(train_data, test_data, D, max_epoch, tag) 
rand('state',0);
randn('state',0);

train = train_data.train;
test = test_data.probe;

N = train_data.N;
M = train_data.M;
K = train_data.K;
mean_rating = train_data.mean_rating; 

num_ratings = size(train,1);   

% SGD parameters
epsilon = 50;
lambda = 0.02;
momentum = 0.8;
num_batches =10;
batch_size = floor(num_ratings/num_batches);

P = 0.1*randn(D,N);
Q = 0.1*randn(D,M);

W = 0.1*randn(D,M);
X = 0.1*randn(D,K);

Y = 0.1*randn(D,K);
Z = 0.1*randn(D,N);

IP = zeros(D,N);
IQ = zeros(D,M);
IW = zeros(D,M);
IX = zeros(D,K);
IY = zeros(D,K);
IZ = zeros(D,N);

err =zeros(max_epoch,1);

for epoch=1:max_epoch
   rr = randperm(num_ratings);	
   train = train(rr,:);
   clear rr;
	
	for batch = 1:num_batches
		fprintf(1,'epoch %d batch %d \r',epoch,batch);

		ii = train((batch-1)*batch_size+1:batch*batch_size,3); 
		jj = train((batch-1)*batch_size+1:batch*batch_size,2); 
		kk = train((batch-1)*batch_size+1:batch*batch_size,1); 
      
		ratings = train((batch-1)*batch_size+1:batch*batch_size,4);
		ratings = double(ratings) - mean_rating;
		
		pred_out = sum(P(:,ii).*Q(:,jj),1)' +  sum(W(:,jj).*X(:,kk),1)' + sum(Y(:,kk).*Z(:,ii),1)';	
		
		R0 = repmat(2*(pred_out-ratings), 1, D)';
		
		DP = R0.*Q(:,jj) + lambda*P(:,ii);
		DQ = R0.*P(:,ii) + lambda*Q(:,jj);
		
		DW = R0.*X(:,kk) + lambda*W(:,jj);
		%DX = R0.*W(:,jj) + lambda*X(:,kk);
		DX(:,kk>1) = R0(:,kk>1).*W(:,jj(kk>1)) + lambda*(X(:,kk(kk>1))-X(:,kk(kk>1)-1));
		DX(:,kk==1) = R0(:,kk==1).*W(:,jj(kk==1)) + lambda*X(:,kk(kk==1));	
		
		DY(:,kk>1) = R0(:,kk>1).*Z(:,ii(kk>1)) + lambda*(Y(:,kk(kk>1))-Y(:,kk(kk>1)-1));
		DY(:,kk==1) = R0(:,kk==1).*Z(:,ii(kk==1)) + lambda*Y(:,kk(kk==1));
		%DY = R0.*Z(:,ii) + lambda*Y(:,kk);	
		DZ = R0.*Y(:,kk) + lambda*Z(:,ii);	
		
		SP = zeros(D,N);		
		SQ = zeros(D,M);		
		SW = zeros(D,M);		
		SX = zeros(D,K);		
		SY = zeros(D,K);		
		SZ = zeros(D,N);		
		
		for it=1:batch_size
     	   SP(:,ii(it)) =  SP(:,ii(it)) +  DP(:,it);
     	   SQ(:,jj(it)) =  SQ(:,jj(it)) +  DQ(:,it);
     	   SW(:,jj(it)) =  SW(:,jj(it)) +  DW(:,it);
     	   SX(:,kk(it)) =  SX(:,kk(it)) +  DX(:,it);
     	   SY(:,kk(it)) =  SY(:,kk(it)) +  DY(:,it);
     	   SZ(:,ii(it)) =  SZ(:,ii(it)) +  DZ(:,it);
		end
	
		IP = momentum*IP + epsilon*SP/batch_size;
		IQ = momentum*IQ + epsilon*SQ/batch_size;
		IW = momentum*IW + epsilon*SW/batch_size;
		IX = momentum*IX + epsilon*SX/batch_size;
		IY = momentum*IY + epsilon*SY/batch_size;
		IZ = momentum*IZ + epsilon*SZ/batch_size;
		
		P =  P - IP;
		Q =  Q - IQ;
		W =  W - IW;
		X =  X - IX;
		Y =  Y - IY;
		Z =  Z - IZ;
	end

	out_ratings = predict(test, P, Q, W, X, Y, Z);
	out_ratings = out_ratings + mean_rating; 
	out_ratings(out_ratings>5) = 5;     
	out_ratings(out_ratings<1) = 1;     
	rmse = sqrt(mean((out_ratings-double(test(:,4))).^2));   
	err(epoch) = rmse;
	fprintf(1,'\nepoch %d, test RMSE %2.4f \n',epoch,rmse);
end

PTF2_Wt.P = P;
PTF2_Wt.Q = Q;
PTF2_Wt.W = W;
PTF2_Wt.X = X;
PTF2_Wt.Y = Y;
PTF2_Wt.Z = Z;
PTF2_Wt.err = err;

outfile = sprintf('%s_PTF2_D%d.mat', tag, D);
save(outfile, 'PTF2_Wt');	
